import psycopg2
import sys
import os

# 当前文件的路径
cur_path = os.path.abspath(os.path.dirname(__file__))
# 上一层目录
parent_path = os.path.dirname(cur_path)
sys.path.append(parent_path)
import config  # 导入配置模块


def connect_to_db():
    """连接 PostgreSQL 数据库"""
    try:
        # 建立连接
        conn = psycopg2.connect(
            host=config.PG_HOST,
            port=config.PG_PORT,
            dbname=config.DBNAME,
            user=config.USER,
            password=config.PASSWORD,
            options=f'-c search_path={config.SCHEMA}'  # 设置默认模式
        )

        print("数据库连接成功！当前模式为:", config.SCHEMA)
        return conn

    except Exception as e:
        print(f"连接数据库失败: {e}")
        return None, None

# 示例用法
if __name__ == "__main__":
    # 执行连接
    conn = connect_to_db()
    # 创建游标对象
    cursor = conn.cursor()

    if conn:
        # 执行测试查询
        cursor.execute("SELECT * FROM tb_community_name_correlation LIMIT 5")
        result = cursor.fetchall()
        print("测试查询结果:", result)

        # 关闭连接
        cursor.close()
        conn.close()